import numpy as np
initial = [0.7, 0.3, 0., 0., 0., 0.]
id2state = {0: 'Subject', 1: 'Adjective', 2: 'Adverb', 3: 'Verb', 4: 'Object', 5: '<EOS>'}
state2id = {id2state[s]: s for s in id2state}
transition = np.array([[0., 0., 0.3, 0.7, 0., 0.],
                       [0.4, 0.1, 0., 0., 0.5, 0.],
                       [0., 0.3, 0., 0.7, 0., 0.],
                       [0., 0.3, 0.2, 0., 0.5, 0.],
                       [0., 0., 0., 0., 0., 1.],
                      ])
# Sample the initial state of the Markov chain
num_states = len(initial)
h = np.random.choice(num_states, p=initial) #参数p即随机选取对象的各元素概率，实际上第二个参数可以是选取的次数
# Sample subsequent states using the conditional probability table
h_total = []
h_total.append(h)
while(h != state2id['<EOS>']):
    h = np.random.choice(num_states, p=transition[h])
    h_total.append(h)

emission = np.array([
    [0.2, 0.2, 0.2, 0.2, 0, 0, 0, 0, 0, 0, 0, 0, 0.1, 0.1, 0.],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0.5, 0.25, 0.25, 0, 0, 0.],
    [0, 0, 0, 0, 0, 0, 0, 0.25, 0.25, 0.5, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0.3, 0.4, 0.3, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0.2, 0.2, 0, 0, 0, 0, 0, 0, 0, 0, 0.3, 0.3, 0],
    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1.]
    ])
word2id = {'I': 0, 'He': 1, 'Jack': 2, 'Mary': 3, 'likes': 4, 'loves': 5, 'hates': 6, 'really': 7, 'extremely': 8, 'pretty': 9, 'cute': 10, 'adorable': 11, 'cats': 12, 'dogs': 13, '.': 14}
id2word = {word2id[w]: w for w in word2id}

def generate(initial, transition, emission):
    # Sample the initial state
    num_states = len(initial)
    num_words = emission.shape[1]

    # Sample the subsequent states and words
    h = np.random.choice(num_states, p=initial)
    h_total = [h]
    v_total = []
    while(h != state2id['<EOS>']):
        v = np.random.choice(num_words, p=emission[h])
        v_total.append(v)
        h = np.random.choice(num_states, p=transition[h])
        h_total.append(h)
    v = np.random.choice(num_words, p=emission[h])
    v_total.append(v)
    return h_total, v_total

h_total, v_total = generate(initial, transition, emission)
print(h_total)
print(' '.join([id2state[h] for h in h_total]))
print(v_total)
print(' '.join([id2word[v] for v in v_total]))

if __name__ == '__main__':
    generate(initial, transition, emission)